from __future__ import division

import sys, logging, argparse

import numpy, scipy.stats

import common
from retraction_unfolding_filter import RetractionPeak, RetractionCurve, CurveStretch

logger = logging.getLogger("curve_stretch_outliers")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

# TODO: Move to a common file and import from here and from kinetics_opt.
bestfloat = numpy.longfloat

def trim_curve_stretch(curve_stretch, s, e, keep_substretches):
    i = -1
    n = len(curve_stretch)
    while i + 1 < n and s <= curve_stretch.peak_for_peaks_ahead(
        i + 1).rupture_force_n <= e:
        i += 1
    if i + 1 != n and not keep_substretches:
        print(i, curve_stretch)
        i = -1
    return curve_stretch.keep_only_last_peaks(i + 1)

def filter_curve_stretches(
    in_curve_stretch_playlist, outlier_test, force_low_high_threshold,
    keep_substretches, out_curve_stretch_playlist):
    with open(in_curve_stretch_playlist, "rb") as in_file:
        in_curve_stretches = [eval(line) for line in in_file]
    v2indices = CurveStretch.v2indices(in_curve_stretches)
    v_indices_list = list(v2indices.items())
    forces_list = CurveStretch.forces_list(
        [[in_curve_stretches[index] for index in indices]
         for _, indices in v_indices_list])
    s_e_list = [
        outlier_bounds(forces, outlier_test, force_low_high_threshold)
        for forces in forces_list]
    filtered_curve_stretches_list = [
        [trim_curve_stretch(in_curve_stretches[index], s, e, keep_substretches)
         for index in indices]
        for (v, indices), (s, e) in zip(v_indices_list, s_e_list)]
    out_curve_stretches = [None] * len(in_curve_stretches)
    for (v, indices), filtered_curve_stretches in zip(
        v_indices_list, filtered_curve_stretches_list):
        for index, filtered_curve_stretch in zip(indices, filtered_curve_stretches):
            common.check_None(out_curve_stretches[index])
            common.check_equal(
                v, in_curve_stretches[index].retraction_velocity_m_per_s)
            common.check_equal(v, filtered_curve_stretch.retraction_velocity_m_per_s)
            out_curve_stretches[index] = filtered_curve_stretch
    with open(out_curve_stretch_playlist, "wb") as out_file:
        for curve_stretch in out_curve_stretches:
            out_file.write(repr(curve_stretch) + "\n")
    logger.info(
        "(Velocity (nm/s), Before, After): %s" % ", ".join(
            ["(%.1f, %d, %d)" % (
                    v * 1e9, len(forces),
                    sum(map(len, filtered_curve_stretches)))
             for (v, _), forces, filtered_curve_stretches in zip(
                    v_indices_list, forces_list, filtered_curve_stretches_list)]))

def outlier_bounds(f, outlier_test, force_low_high_threshold):
    if not len(f):
        return 0, numpy.inf
    if outlier_test == "iqr":
        f = numpy.asarray(f, dtype=bestfloat)
        q1, q3 = scipy.stats.scoreatpercentile(f, 25), scipy.stats.scoreatpercentile(f, 75)
        s, e = q1 - force_low_high_threshold * (q3 - q1), q3 + force_low_high_threshold * (q3 - q1)
    elif outlier_test == "upper_iqr":
        f = numpy.asarray(f, dtype=bestfloat)
        q1, q3 = scipy.stats.scoreatpercentile(f, 25), scipy.stats.scoreatpercentile(f, 75)
        s, e = 0, q3 + force_low_high_threshold * (q3 - q1)
    elif outlier_test == "iqr2":
        f = numpy.asarray(f, dtype=bestfloat)
        q1, q3 = scipy.stats.scoreatpercentile(f, 25), scipy.stats.scoreatpercentile(f, 75)
        s, e = q1 - 2 * force_low_high_threshold * (q3 - q1), q3 + force_low_high_threshold * (q3 - q1)
    elif outlier_test == "median_z":
        f_median = numpy.median(f)
        f_mad = numpy.median(abs(f - f_median))
        k = 1 / scipy.stats.norm.ppf(.75)  # about 1.4826
        f_std = k * f_mad
        s, e = f_median - force_low_high_threshold * f_std, f_median + force_low_high_threshold * f_std
    elif outlier_test == "negative_skew":
        sorted_f = sorted(f)
        j = len(sorted_f)
        while j >= 2:
            skew = scipy.stats.skew(sorted_f[:j])
            logger.info("j=%d skew=%s", j, skew)
            if skew >= 0:
                j -= 1
            else:
                break
        s, e = 0, sorted_f[j - 1]
    elif outlier_test == "min_range":
        #g = scipy.stats.norm(loc=0, sigma=1)
        #fwhm = g.cdf(numpy.sqrt(2 * numpy.log(2))) - g.cdf(-2 * numpy.log(2))  # 0.79765553527354549
        common.check_le(force_low_high_threshold, 1)
        sorted_f = sorted(f)
        n = len(sorted_f)
        m = int(force_low_high_threshold * n)
        i = numpy.argmin(
            [sorted_f[i + m - 1] - sorted_f[i] for i in range(n - m + 1)])
        s, e = sorted_f[i], sorted_f[i + m - 1]
    elif outlier_test == "none":
        s, e = min(f), max(f)
    else:
        raise Exception("Unknown outlier test: %s." % outlier_test)
    logger.info("s(pN)=%s e(pN)=%s", s * 1e12, e * 1e12)
    return s, e

def main(params):
    parser = argparse.ArgumentParser()
    parser.add_argument("--in_curve_stretch_playlist")
    parser.add_argument(
        "--outlier_test", choices=(
            "none", "upper_iqr", "iqr", "iqr2", "z", "median_z", "negative_skew",
            "min_range"),
        help="iqr is the Inter-Quartile Range")
    parser.add_argument(
        "--force_low_high_threshold", type=float,
        help="mathforum.org/library/drmath/view/52720.html")
    parser.add_argument(
        "--keep_substretches", choices=("True", "true", "False", "false"))
    parser.add_argument("--out_curve_stretch_playlist")
    
    args = parser.parse_args(params)
    logger.info("%s", args)

    keep_substretches = {"false" : False, "true" : True}.get(
        args.keep_substretches.lower())

    filter_curve_stretches(
        args.in_curve_stretch_playlist, args.outlier_test,
        args.force_low_high_threshold, keep_substretches,
        args.out_curve_stretch_playlist)


if __name__ == "__main__":
    main(sys.argv[1:])
